{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Initialize\n",
    "from ochumanApi.ochuman import OCHuman\n",
    "# <Filter>: \n",
    "#      None(default): load all. each has a bbox. some instances have keypoint and some have mask annotations.\n",
    "#            images: 5081, instances: 13360\n",
    "#     'kpt&segm' or 'segm&kpt': only load instances contained both keypoint and mask annotations (and bbox)\n",
    "#            images: 4731, instances: 8110\n",
    "#     'kpt|segm' or 'segm|kpt': load instances contained either keypoint or mask annotations (and bbox)\n",
    "#            images: 5081, instances: 10375\n",
    "#     'kpt' or 'segm': load instances contained particular kind of annotations (and bbox)\n",
    "#            images: 5081/4731, instances: 10375/8110\n",
    "ochuman = OCHuman(AnnoFile='./ochuman.json', Filter='kpt&segm')\n",
    "image_ids = ochuman.getImgIds()\n",
    "print ('Total images: %d'%len(image_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load annotations & Visualize\n",
    "import cv2, os\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (15, 15)\n",
    "import ochumanApi.vis as vistool\n",
    "from ochumanApi.ochuman import Poly2Mask\n",
    "\n",
    "ImgDir = './images/'\n",
    "data = ochuman.loadImgs(imgIds=[image_ids[2]])[0]\n",
    "\n",
    "img = cv2.imread(os.path.join(ImgDir, data['file_name']))\n",
    "height, width = data['height'], data['width']\n",
    "\n",
    "colors = [[255, 0, 0], \n",
    "         [255, 255, 0],\n",
    "         [0, 255, 0],\n",
    "         [0, 255, 255], \n",
    "         [0, 0, 255], \n",
    "         [255, 0, 255]]\n",
    "\n",
    "for i, anno in enumerate(data['annotations']):\n",
    "    bbox = anno['bbox']\n",
    "    kpt = anno['keypoints']\n",
    "    segm = anno['segms']\n",
    "    max_iou = anno['max_iou']\n",
    "\n",
    "    img = vistool.draw_bbox(img, bbox, thickness=3, color=colors[i%len(colors)])\n",
    "    if segm is not None:\n",
    "        mask = Poly2Mask(segm)\n",
    "        img = vistool.draw_mask(img, mask, thickness=3, color=colors[i%len(colors)])\n",
    "    if kpt is not None:\n",
    "        img = vistool.draw_skeleton(img, kpt, connection=None, colors=colors[i%len(colors)], bbox=bbox)\n",
    "\n",
    "plt.imshow(img[:,:,::-1])\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Convert to coco format.\n",
    "# <maxIouRange>: \n",
    "#     (0.0, 1.0) means all instances\n",
    "#     (0.5, 0.75) means Moderate instances\n",
    "#     (0.75, 1.0) means Hard instances\n",
    "ochuman.toCocoFormart(subset='val', maxIouRange=(0.0, 1.0), save_dir='./')\n",
    "ochuman.toCocoFormart(subset='test', maxIouRange=(0.0, 1.0), save_dir='./')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
